#include "common.h"
#include <string.h>

void *srcbuf_h = NULL;
void *srcbuf_d = NULL;
void *dstbuf_h = NULL;
void *dstbuf_d = NULL;
void *relaybuf_h = NULL;

int buffer_init(size_t size) {
  srcbuf_h = malloc(size);
  if (srcbuf_h == NULL) {
    fprintf(stderr, "malloc srcbuf_h failed\n");
    return -1;
  }
  dstbuf_h = malloc(size);
  if (dstbuf_h == NULL) {
    fprintf(stderr, "malloc dstbuf_h failed\n");
    free(srcbuf_h);
    return -1;
  }
  relaybuf_h = malloc(size);
  if (relaybuf_h == NULL) {
    fprintf(stderr, "malloc relaybuf_h failed\n");
    free(srcbuf_h);
    free(dstbuf_h);
    return -1;
  }
  for (size_t i = 0; i < size; i++) {
    ((char *)srcbuf_h)[i] = (char)(i % 256);
  }
  ACL_CHECK(aclrtSetDevice(FROM_DEVICE));
  ACL_CHECK(aclrtMalloc(&srcbuf_d, size, ACL_MEM_MALLOC_HUGE_FIRST_P2P));
  ACL_CHECK(aclrtSetDevice(TO_DEVICE));
  ACL_CHECK(aclrtMalloc(&dstbuf_d, size, ACL_MEM_MALLOC_HUGE_FIRST_P2P));
  ACL_CHECK(aclrtSetDevice(FROM_DEVICE));
  ACL_CHECK(
      aclrtMemcpy(srcbuf_d, size, srcbuf_h, size, ACL_MEMCPY_HOST_TO_DEVICE));
  return 0;
}

int buffer_destroy() {
  ACL_CHECK(aclrtSetDevice(FROM_DEVICE));
  ACL_CHECK(aclrtFree(srcbuf_d));
  ACL_CHECK(aclrtSetDevice(TO_DEVICE));
  ACL_CHECK(aclrtFree(dstbuf_d));
  free(srcbuf_h);
  free(dstbuf_h);
  free(relaybuf_h);
  return 0;
}

int buffer_verify() {
  ACL_CHECK(aclrtSetDevice(TO_DEVICE));
  ACL_CHECK(aclrtMemcpy(dstbuf_h, DATA_SIZE, dstbuf_d, DATA_SIZE,
                        ACL_MEMCPY_DEVICE_TO_HOST));
  if (memcmp(srcbuf_h, dstbuf_h, DATA_SIZE) != 0) {
    fprintf(stderr, "data verify failed\n");
    return -1;
  }
  return 0;
}

int buffer_verify_half() {
  ACL_CHECK(aclrtSetDevice(TO_DEVICE));
  ACL_CHECK(aclrtMemcpy(dstbuf_h, DATA_SIZE, dstbuf_d, DATA_SIZE,
                        ACL_MEMCPY_DEVICE_TO_HOST));
  if (memcmp(srcbuf_h, dstbuf_h, DATA_SIZE / 2) != 0) {
    fprintf(stderr, "first half data verify failed\n");
    return -1;
  }

  if (memcmp((char *)srcbuf_h + DATA_SIZE / 2,
             (char *)relaybuf_h + DATA_SIZE / 2, DATA_SIZE / 2) != 0) {
    fprintf(stderr, "second half data verify failed\n");
    return -1;
  }
  return 0;
}